package com.atlinmengze.java.ai.langchain4j.tools;

import dev.langchain4j.agent.tool.Tool;
import org.springframework.stereotype.Component;
import java.util.*;

@Component
public class DepressionAssessmentTools {
    // 内存存储用户答题进度，实际可用MongoDB等持久化
    private final Map<Long, AssessmentSession> sessionMap = new HashMap<>();

    // PHQ-9 题库
    private static final List<Question> PHQ9_QUESTIONS = List.of(
            new Question("做事时提不起兴趣或没有乐趣", List.of("完全不会", "有几天", "一半以上的天数", "几乎每天")),
            new Question("感到心情低落、沮丧或绝望", List.of("完全不会", "有几天", "一半以上的天数", "几乎每天")),
            new Question("入睡困难、易醒或睡眠过多", List.of("完全不会", "有几天", "一半以上的天数", "几乎每天")),
            new Question("感觉疲倦或没有精力", List.of("完全不会", "有几天", "一半以上的天数", "几乎每天")),
            new Question("食欲不振或暴饮暴食", List.of("完全不会", "有几天", "一半以上的天数", "几乎每天")),
            new Question("觉得自己很糟或失败，或让家人失望", List.of("完全不会", "有几天", "一半以上的天数", "几乎每天")),
            new Question("注意力难以集中（如阅读或看电视时）", List.of("完全不会", "有几天", "一半以上的天数", "几乎每天")),
            new Question("动作或说话缓慢到别人能察觉，或相反（烦躁、坐立不安）", List.of("完全不会", "有几天", "一半以上的天数", "几乎每天")),
            new Question("有不如死掉或用某种方式伤害自己的念头", List.of("完全不会", "有几天", "一半以上的天数", "几乎每天"))
    );

    private static final int[][] PHQ9_SCORE_RANGES = {
            {0, 4},   // 无明显抑郁症状
            {5, 9},   // 轻度抑郁
            {10, 14}, // 中度抑郁
            {15, 19}, // 中重度抑郁
            {20, 27}  // 重度抑郁
    };
    private static final String[] PHQ9_LEVELS = {
            "无明显抑郁症状", "轻度抑郁", "中度抑郁", "中重度抑郁", "重度抑郁"
    };

    @Tool(name = "PHQ-9抑郁自评", value = "引导用户完成PHQ-9量表自评，自动评分并输出结果")
    public String startPHQ9(Long memoryId) {
        AssessmentSession session = new AssessmentSession(PHQ9_QUESTIONS);
        sessionMap.put(memoryId, session);
        return askNextQuestion(memoryId, "欢迎使用PHQ-9抑郁自评量表。请如实作答。\n");
    }

    // 用户答题后调用
    public String answerPHQ9(Long memoryId, int optionIndex) {
        AssessmentSession session = sessionMap.get(memoryId);
        if (session == null) return "请先输入\"我要做抑郁自评\"开始测试。";
        session.answer(optionIndex);
        if (session.isFinished()) {
            int score = session.getTotalScore();
            String level = getLevel(score);
            sessionMap.remove(memoryId);
            StringBuilder sb = new StringBuilder();
            sb.append("自评完成，总分：").append(score).append("，抑郁程度：").append(level).append("。\n");
            if ("重度抑郁".equals(level)) {
                sb.append("建议您尽快寻求专业医生帮助，及时就医！");
            }
            return sb.toString();
        } else {
            return askNextQuestion(memoryId, null);
        }
    }

    // 输出下一题
    private String askNextQuestion(Long memoryId, String prefix) {
        AssessmentSession session = sessionMap.get(memoryId);
        Question q = session.getCurrentQuestion();
        StringBuilder sb = new StringBuilder();
        if (prefix != null) sb.append(prefix);
        sb.append("第").append(session.getCurrentIndex() + 1).append("题：").append(q.text).append("\n");
        for (int i = 0; i < q.options.size(); i++) {
            sb.append(i + 1).append(". ").append(q.options.get(i)).append("\n");
        }
        sb.append("请回复选项序号（1-4）。");
        return sb.toString();
    }

    private String getLevel(int score) {
        for (int i = 0; i < PHQ9_SCORE_RANGES.length; i++) {
            int[] range = PHQ9_SCORE_RANGES[i];
            if (score >= range[0] && score <= range[1]) return PHQ9_LEVELS[i];
        }
        return "未知";
    }

    // 内部类：答题会话
    static class AssessmentSession {
        private final List<Question> questions;
        private int currentIndex = 0;
        private final List<Integer> answers = new ArrayList<>();
        public AssessmentSession(List<Question> questions) { this.questions = questions; }
        public void answer(int optionIndex) { answers.add(optionIndex - 1); currentIndex++; }
        public boolean isFinished() { return currentIndex >= questions.size(); }
        public Question getCurrentQuestion() { return questions.get(currentIndex); }
        public int getCurrentIndex() { return currentIndex; }
        public int getTotalScore() {
            int sum = 0;
            for (int i = 0; i < answers.size(); i++) {
                sum += answers.get(i);
            }
            return sum;
        }
    }

    // 内部类：题目
    static class Question {
        String text;
        List<String> options;
        public Question(String text, List<String> options) {
            this.text = text;
            this.options = options;
        }
    }
} 